package com.wonder.fileSpit;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

/**
 * 多线程文件拆分工具类。
 * 默认功能：根据文本行号求模算法，将不同行的文本内容，写入不同文件。
 * 增强功能：按照列值hash后取模算法，将相同字段值的行写入同一个文件，保证有序。（例如不同部门的csv数据，可以按部门字段拆分，相同部门放到同一个文件）
 * 适用于将一个大文件拆分成多个文件的工具类
 * @author: jinrong.wang
 * @date: 2020/12/23
 */
public class FileSpitUtil {




    /**
     * 结束标记，最后一行写写入后，再写一个结束标记
     */
    private static final String FILE_WRITE_OVER = "FILE_WRITE_OVER";

    /**
     * 如果spitIndex和spitSymbol都为空，则按行数与文件数取模，将数据写入对应的线程
     * 如果spitIndex和spitSymbol都非空，则按列值hash取模算法，将相同字段值的行写入同一个文件，保证有序。
     *
     * @param sourcePath 源文件路径
     * @param targetPath 目标文件路径
     * @param fileName 文件名称
     * @param fileNums 文件数量
     * @param isHaveHeader 是否有表头,true则每一个文件都写一份表头，false则跳过表头
     * @param spitIndex 分割字段索引
     * @param spitSymbol 分割符号
     * @return
     */
    public List<String> fileSpit(String sourcePath, String targetPath, String fileName, int fileNums, boolean isHaveHeader, Integer spitIndex, String spitSymbol) {
        //开始时间
        long startTime = System.currentTimeMillis();
        //文件列表
        List<String> fileNameList = new ArrayList<>();
        //判断文件是否全部拆分成功
        boolean isAllSuccess = true;
        //线程列表
        Map<Integer, FileWriteCallable> map = new HashMap<>();
        //线程结果列表
        Map<Integer, Future<Boolean>> mapFuture = new HashMap<>();
        //线程池
        ExecutorService executorService = Executors.newFixedThreadPool(fileNums);
        //拆成几个文件，则启动几个异步线程
        for (int i = 0; i < fileNums; i++) {
            FileWriteCallable fileWriteCallable = new FileWriteCallable();
            //目标文件名：/tmp/0.test.txt
            String filePath = targetPath + i + "." + fileName;
            //初始化线程参数
            fileWriteCallable.init(FILE_WRITE_OVER, filePath, 1000);
            //提交线程
            Future<Boolean> future = executorService.submit(fileWriteCallable);
            //缓存线程信息，方便后续关闭
            map.put(i, fileWriteCallable);
            mapFuture.put(i, future);
            fileNameList.add(filePath);
        }

        //开始解析源文件，并将每一行数据，按照拆分算法放到不同文件
        try (InputStream inputStream = new FileInputStream(sourcePath + fileName);
             InputStreamReader inputStreamReader = new InputStreamReader(inputStream);
             BufferedReader reader = new BufferedReader(inputStreamReader)
        ) {
            int lineNumber = 0;
            String line;
            while ((line = reader.readLine()) != null) {

                //第一行根据isHaveHeader判断是否在每一个文件都写入，例如csv有表头，拆分后可以在每个文件都写入表头
                if (lineNumber == 0) {
                    if (isHaveHeader) {
                        for (int i = 0; i < fileNums; i++) {
                            map.get(i).add(line);
                        }
                        lineNumber++;
                        continue;
                    } else {
                        System.out.println("不需要表头");
                    }
                }
                //拆分算法
                try {
                    //如果spitIndex和spitSymbol都为空，则按行数与文件数取模，将数据写入对应的线程
                    if (spitIndex == null || spitSymbol == null) {
                        //按行数与文件数取模，将数据写入对应的线程
                        int fileIndex = lineNumber % fileNums;
                        map.get(fileIndex).add(line);
                    } else {
                        //如果spitIndex和spitSymbol都非空，则按列值hash取模算法，将相同字段值的行写入同一个文件，保证有序。
                        //根据某一列值，hash值取模，将数据写入对应的线程。例如根据部门编号拆分，将部门编号相同的数据写入同一个文件
                        //将一行数据按分割符进行拆分
                        String[] csvLine = line.split(spitSymbol);
                        //将列值作为key
                        String key = csvLine[spitIndex];
                        //hash值与拆分文件数求模，Integer.MAX_VALUE位运算避免出现负数
                        int fileIndex = (key.hashCode() & Integer.MAX_VALUE) % fileNums;
                        //写入对应的线程
                        map.get(fileIndex).add(line);
                    }
                } catch (Exception e) {
                    //无法拆分，写入第0个线程
                    map.get(0).add(line);
                    System.out.println("拆分异常文件，写入第0个线程");
                    e.printStackTrace();
                }
                lineNumber++;
            }
            System.out.println("文件解析fileName=" + fileName + ",lineNumber=" + lineNumber);
        } catch (Exception e) {
            e.printStackTrace();
            isAllSuccess = false;
        }

        //源文件解析完成，但是线程可能没有执行完，或者IO流还有缓存。需要再写入一个结束标记，等待线程线程执行完成。
        for (int i = 0; i < fileNums; i++) {
            //添加结束标记，线程会自动结束
            map.get(i).add(FILE_WRITE_OVER);
            //get方法阻塞主线程，等待拆分线程执行完成
            Future<Boolean> future = mapFuture.get(i);
            boolean flag = false;
            try {
                //获取线程写入文件结果，如果返回false，则说明线程有io异常，或者其他异常
                flag = future.get(3, TimeUnit.SECONDS);
            } catch (Exception e) {
                e.printStackTrace();
                isAllSuccess = false;
            }
            if (flag) {
                System.out.println("线程已关停");
            } else {
                System.out.println("线程关停失败");
                isAllSuccess = false;
            }
        }
        System.out.println("耗时(秒)=" + (System.currentTimeMillis() - startTime) / 1000);
        //强制关停线程。
        executorService.shutdownNow();
        if (!isAllSuccess) {
            throw new RuntimeException("文件分割不全成功");
        } else {
            return fileNameList;
        }
    }

    public static void main(String[] args) {
        FileSpitUtil fileSpitUtil = new FileSpitUtil();
        //按第0个字段hash值与文件数3取余数，将数据写入对应的线程
        fileSpitUtil.fileSpit("D:\\temp\\", "D:\\temp\\", "新建文本文档.txt",
                3, false,
                0, ",");

        //按行数与文件数3取模，将数据写入对应的线程
        fileSpitUtil.fileSpit("D:\\temp\\", "D:\\temp\\", "新建文本文档2.txt",
                3, false,
                null, null);
    }
}
